import os
from keras import Sequential, optimizers, losses
from keras.layers import Dense
from tensorflow import saved_model

from prepare_mnist_data import load_mnist_data, process_mnist_data

if __name__ == '__main__':
    os.environ['TF_CPP_MIN_LOG_LEVEl'] = '2'
    (train_images, train_labels), (test_images, test_labels) = load_mnist_data()

    db, ds_val = process_mnist_data(train_images, train_labels, test_images, test_labels, 128)

    model = Sequential(
        [
            Dense(units=256,activation='relu'),
            Dense(units=128,activation='relu'),
            Dense(units=64,activation='relu'),
            Dense(units=32,activation='relu'),
            Dense(units=10),
        ]
    )

    model.build(input_shape=(None, 28*28))
    model.summary()

    model.compile(optimizer=optimizers.Adam(learning_rate=0.001),
                  loss=losses.CategoricalCrossentropy(from_logits=True),
                  metrics=['Accuracy'])


    model.fit(db, epochs=3,validation_data=ds_val, validation_freq=2)

    model.evaluate(ds_val)

    # model.save("mnist_model.h5")
    saved_model.save(model, "mnist_savemodel")
